import json, yaml, argparse,os
from multiprocessing import freeze_support
import pytorch_lightning as pl
from openue.lit_models import RELitModel
from openue.data import REDataset

with open("ske/train.json", "r",encoding='utf8') as file:
    for line in file.readlines():
        example = json.loads(line)
        break
for k, v in example.items():
    print(f"{k}: {v}")


early_callback = pl.callbacks.EarlyStopping(monitor="Eval/f1", mode="max", patience=5)
model_checkpoint = pl.callbacks.ModelCheckpoint(monitor="Eval/f1", mode="max",
    filename='{epoch}-{Eval/f1:.2f}',
    dirpath="output",
    save_weights_only=True
)

path = "./run_seq.yaml"
# 使用config.yaml 载入超参设置

parser = argparse.ArgumentParser(add_help=False)
args = parser.parse_args(args=[])
opt = yaml.load(open(path))
args.__dict__.update(opt)


callbacks = [early_callback, model_checkpoint]
trainer = pl.Trainer.from_argparse_args(args, callbacks=callbacks, default_root_dir="training/logs")


data = REDataset(args)

lit_model = RELitModel(args, data_config=data.get_config())

def _save_model(litmodel, tokenizer, path):
    os.system(f"mkdir -p {path}")
    litmodel.model.save_pretrained(path)
    tokenizer.save_pretrained(path)

def main():
    trainer.fit(lit_model, datamodule=data)
    trainer.test(lit_model, datamodule=data)

    _save_model(litmodel=lit_model, tokenizer=data.tokenizer, path="seq_model")


if __name__ == '__main__':
    freeze_support()
    main()